import json
import os
import random
import subprocess
import time
import docker
import argparse

import yaml
import platform

client = docker.from_env()


def is_arm64():
    if 'arm64' in platform.machine().lower():
        return True

def _exec_cmd(command):
    return subprocess.getstatusoutput(command)
def select_card():
    str_in = input("请输入需挂载的卡ID （如需多张卡请用逗号隔开）：")
    str_in = str_in.strip()
    res = []
    for it in str_in.split(','):
        print(it)
        res.append(int(it))
    return list(set(res))

def is_npu():
    exit_code, _ =  _exec_cmd('npu-smi')
    return exit_code == 0

def is_gpu():
    exit_code, _ =  _exec_cmd('nvidia-smi')
    return exit_code == 0

def run_cmd(cmd):
    print(f'[exec] {cmd}')
    exit_code, output = _exec_cmd(cmd)
    print(output)
    if exit_code != 0:
        exit(1)
    return output

def load_image(file_path):
    print(f'[load_image] {file_path}')
    with open(file_path, 'rb') as ff:
        images = client.images.load(ff.read())
    image = images[0]

    return image.id


def is_container_running(container_id):
    exit_code, output = subprocess.getstatusoutput(f'docker inspect {container_id}')
    if exit_code != 0:
        return False
    try:
        res = json.loads(output)
        return res[0]['State']['Running']
    except Exception:
        return False


def get_container_logs(container_id, tail=100):
    exit_code, output = subprocess.getstatusoutput(f'docker logs  --tail {tail} {container_id}')
    if exit_code != 0:
        return ''
    return output


def stop_docker(name):
    run_cmd(f'docker rm -f {name}')


if not os.path.exists('LICENSE'):
    print('缺少LICENSE授权')
    exit(0)

if not os.path.exists('config.yaml'):
    print('缺少配置文件config.yaml')
    exit(0)

with open('config.yaml', 'r') as f:
    content = f.read()

with open('LICENSE', 'r') as f:
    license = f.read()
license = license.strip()

cfg = yaml.load(content, yaml.FullLoader)

if is_arm64():
    license_server_image = load_image("license-server-arm64.tar")
else:
    license_server_image = load_image("license-server-amd64.tar")

for k, v in cfg.items():
    image_file = v.get('ImageFile')
    if image_file and len(image_file) > 0:
        load_image(image_file)

license_server = cfg['LicenseServer']
stop_docker(license_server["Name"])
c_id = run_cmd(
    f'docker run -d --name {license_server["Name"]}  --network=host --privileged=true -e LICENSE={str(license)} {license_server_image}')
print(f'等待镜像 {license_server["Name"]} 启动中 ......')
time.sleep(10)
if not is_container_running(c_id):
    print('镜像启动失败，日志：')
    print(get_container_logs(c_id))
    exit(1)

model_service = cfg['ModelService']
stop_docker(model_service["Name"])

model_service_image = load_image(model_service["ImageFile"])

volumes = model_service.get('Volumes', [])
port = random.randint(10000, 20000)
cmd = f'docker run --name {model_service["Name"]} -d  -p {port}:8080'
for item in volumes:
    cmd += f' -v {item["HostPath"]}:{item["ContainerPath"]} '

card_ids = select_card()
print(f"挂载的卡ID为：{card_ids}")

if len(card_ids) > 0:
    if is_npu():
        for cid in card_ids:
            cmd += f' --device=/dev/davinci{cid} '

        cmd += f' --ipc=host '
        cmd += f' --device=/dev/davinci_manager '
        cmd += f' --device=/dev/devmm_svm '
        cmd += f' --device=/dev/hisi_hdc '
        cmd += f' -v /usr/local/Ascend/driver:/usr/local/Ascend/driver '
        cmd += f' -v /usr/local/Ascend/add-ons/:/usr/local/Ascend/add-ons/ '

    if is_gpu():
        str_ids = []
        for cid in card_ids:
            str_ids.append(str(cid))
        device_ids = ','.join(str_ids)
        cmd += f' –gpus "device={device_ids}" '

c_id = run_cmd(f'{cmd} {model_service_image}')
print(f'等待镜像 {model_service["Name"]} 启动中 ......')
time.sleep(10)
if not is_container_running(c_id):
    print('镜像启动失败，日志：')
    print(get_container_logs(c_id))
    exit(1)

print(f'模型服务启动成功, http://localhost:{port}')



if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('--workdir', '-w', type=str, default='.', help="workdir")
    parser.add_argument('--add_dir', type=str, action='append', default=[], required=False, help="add path")
    parser.add_argument('--output', '-o', type=str, required=False, default='out.tar')
    parser.add_argument('--add_file', type=str, action='append', default=[], help="add file")
    parser.add_argument('--add_image', type=str, action='append', default=[], help="add image")

    args = parser.parse_args()
